# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/22 09:21 
@Author  : ZhangShenao 
@File    : 1.相似性检索.py
@Desc    : 使用Pinecone进行相似性检索

Pinecone支持原生的带过滤条件的相似性检索功能,即元数据筛选
使用元数据过滤器会精确检索与过滤器匹配的最接近的结果数
过滤器格式为json,其中键是元数据字段对应字符串,值是字段要匹配的值

"""
import os

import dotenv
from langchain_openai import OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore

# 加载环境变量
dotenv.load_dotenv()

# 创建Embeddings
embeddings = OpenAIEmbeddings(
    model='text-embedding-3-small',  # 选择text-embedding-3-small模型
    openai_api_base=os.getenv('OPENAI_API_BASE')
)

# 创建Pinecone向量数据库
vector_store = PineconeVectorStore(
    index_name='llm-ops',  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace='llm-ops',  # 指定Namespace
    text_key='text',  # 指定文本属性key
)

# 带相关性得分的相似性检索
similarity_search_result = vector_store.similarity_search_with_relevance_scores(
    query='我养了一只猫叫笨笨',  # 指定搜索关键词
    k=2,  # 指定返回文档的条数
)
print(f'similarity_search_result: {similarity_search_result}')

# 带过滤条件的相似性检索
similarity_search_with_filter_result = vector_store.similarity_search_with_relevance_scores(
    query='我养了一只猫叫笨笨',  # 指定搜索关键词
    k=4,  # 指定返回文档的条数
    # 指定过滤器,筛选元数据中page大于等于5的文档
    filter={
        'page': {'$gte': 5}
    }
)
print(f'similarity_search_with_filter_result: {similarity_search_with_filter_result}')
